#include "HeadgenAstConsumer.h"

#include <clang/AST/Decl.h>
#include <clang/AST/DeclGroup.h>
#include <clang/AST/PrettyPrinter.h>
#include <clang/AST/Expr.h>
#include <clang/Basic/SourceManager.h>
#include <clang/Basic/FileManager.h>
#include <clang/Basic/LangOptions.h>

#include <llvm/Support/raw_ostream.h>

HeadgenAstConsumer::HeadgenAstConsumer(clang::SourceManager &sm,
									   LinkedIdentifierInfoMap &idMap,
									   LinkedIdentifierInfoList &idList) :
	m_sourceManager(sm),
	m_idMap(idMap),
	m_idList(idList)
{}

HeadgenAstConsumer::~HeadgenAstConsumer() {
}

void HeadgenAstConsumer::HandleTopLevelDecl(clang::DeclGroupRef D) {
	clang::DeclGroupRef::iterator it;
	for (it = D.begin(); it != D.end(); ++ it) {
		clang::FunctionDecl *funcDecl = llvm::dyn_cast_or_null<clang::FunctionDecl>(*it);
		clang::RecordDecl *recordDecl = llvm::dyn_cast_or_null<clang::RecordDecl>(*it);
		clang::TypedefDecl *typedefDecl = llvm::dyn_cast_or_null<clang::TypedefDecl>(*it);
		clang::VarDecl *varDecl = llvm::dyn_cast_or_null<clang::VarDecl>(*it);
		if (funcDecl) {
			processFunctionDecl(funcDecl);
		}
		else if (recordDecl) {
			processRecordDecl(recordDecl);
		}
		else if (typedefDecl) {
			processTypedefDecl(typedefDecl);
		}
		else if (varDecl) {
			processVarDecl(varDecl);
		}
	}
}


void HeadgenAstConsumer::processFunctionDecl(clang::FunctionDecl *D) {
	if (D->isFirstDeclaration()) {
		std::string name = D->getDeclName().getAsString();
		LinkedIdentifierInfo idInfo = LinkedIdentifierInfo(
					name, D, m_sourceManager.getFileEntryForID(
						m_sourceManager.getFileID(D->getLocation()))->getName());

		if (m_sourceManager.isFromMainFile(D->getLocation())) {
			std::list<std::string> deps;
			getTypeDeps(D->getResultType(), deps);
			for (unsigned i = 0, e = D->getNumParams(); i != e; ++i) {
				getTypeDeps(D->getParamDecl(i)->getType(), deps);
			}
			for (std::list<std::string>::iterator i = deps.begin(); i != deps.end(); ++ i) {
				LinkedIdentifierInfoMap::iterator iId = m_idMap.find(*i);
				if (iId != m_idMap.end()) {
					idInfo.m_deps.push_back(&(iId->second));
				}
			}
		}

		m_idMap[name] = idInfo;
		m_idList.push_back(&m_idMap[name]);
	}
}


void HeadgenAstConsumer::processRecordDecl(clang::RecordDecl *D) {
	if (D->isFirstDeclaration()) {
		std::string name = D->getDeclName().getAsString();
		LinkedIdentifierInfo idInfo = LinkedIdentifierInfo(
					name, D, m_sourceManager.getFileEntryForID(
						m_sourceManager.getFileID(D->getLocation()))->getName());

		if (m_sourceManager.isFromMainFile(D->getLocation())) {
			std::list<std::string> deps;
			for (clang::RecordDecl::field_iterator i = D->field_begin();
				 i != D->field_end(); ++i) {
				getTypeDeps((*i)->getType(), deps);
			}
			for (std::list<std::string>::iterator i = deps.begin(); i != deps.end(); ++ i) {
				LinkedIdentifierInfoMap::iterator iId = m_idMap.find(*i);
				if (iId != m_idMap.end()) {
					idInfo.m_deps.push_back(&(iId->second));
				}
			}
		}

		m_idMap[name] = idInfo;
		m_idList.push_back(&m_idMap[name]);
	}
}


void HeadgenAstConsumer::processTypedefDecl(clang::TypedefDecl *D) {
	if (D->isFirstDeclaration()) {
		std::string name = D->getDeclName().getAsString();
		LinkedIdentifierInfo idInfo = LinkedIdentifierInfo(
					name, D, m_sourceManager.getFileEntryForID(
						m_sourceManager.getFileID(D->getLocation()))->getName());

		if (m_sourceManager.isFromMainFile(D->getLocation())) {
			std::list<std::string> deps;
			getTypeDeps(D->getTypeSourceInfo()->getType(), deps);
			for (std::list<std::string>::iterator i = deps.begin(); i != deps.end(); ++ i) {
				LinkedIdentifierInfoMap::iterator iId = m_idMap.find(*i);
				if (iId != m_idMap.end()) {
					idInfo.m_deps.push_back(&(iId->second));
				}
			}
		}

		m_idMap[name] = idInfo;
		m_idList.push_back(&m_idMap[name]);
	}
}


void HeadgenAstConsumer::processVarDecl(clang::VarDecl *D) {
	if (D->isFirstDeclaration()) {
		std::string name = D->getDeclName().getAsString();
		LinkedIdentifierInfo idInfo = LinkedIdentifierInfo(
					name, D, m_sourceManager.getFileEntryForID(
						m_sourceManager.getFileID(D->getLocation()))->getName());

		if (m_sourceManager.isFromMainFile(D->getLocation())) {
			std::list<std::string> deps;
			getTypeDeps(D->getType(), deps);
			for (std::list<std::string>::iterator i = deps.begin(); i != deps.end(); ++ i) {
				LinkedIdentifierInfoMap::iterator iId = m_idMap.find(*i);
				if (iId != m_idMap.end()) {
					idInfo.m_deps.push_back(&(iId->second));
				}
			}
		}

		m_idMap[name] = idInfo;
		m_idList.push_back(&m_idMap[name]);
	}
}

// TODO: enable support for unnamed records.
void HeadgenAstConsumer::getTypeDeps(clang::QualType T, std::list<std::string> &res) {
	if (const clang::TypedefType *tt = T.getTypePtr()->getAs<clang::TypedefType>()) {
		res.push_back(tt->getDecl()->getName());
	}
	else if (const clang::RecordType *rt = T.getTypePtr()->getAsStructureType()) {
		res.push_back(rt->getDecl()->getNameAsString());
	}
	else if (const clang::RecordType *rt = T.getTypePtr()->getAsUnionType()) {
		res.push_back(rt->getDecl()->getNameAsString());
	}
	else if (const clang::PointerType *pt = T.getTypePtr()->getAs<clang::PointerType>()) {
		getTypeDeps(pt->getPointeeType(), res);
	}
/*	else if (const clang::ArrayType *at = T.getTypePtr()->getAsArrayTypeUnsafe()) {
		getTypeDeps(at->getElementType(), res);
	}*/
	else {
		res.push_back(T.getUnqualifiedType().getAsString(
						  clang::PrintingPolicy(clang::LangOptions())));
	}
}
